from numpy import dot
from scipy import spatial
from model.MyOllamaEmbeddings import MyOllamaEmbeddings
# pip install numpy
from numpy.linalg import norm

embeddings = MyOllamaEmbeddings(model="nomic-embed-text")


def get_embeddings(text):
    response = embeddings.embed_query(text)
    return response

# 余弦相似度
def cosine_similarity(vector1, vector2):
    # 计算并返回：两个向量之间的相似度，公式：两个向量的点积除以他们范数的乘积
    return dot(vector1, vector2) / (norm(vector1) * norm(vector2))


# 文本搜索功能
# query：查询字符串，
# documents：文档列表
def search_documents(query, documents):
    # 1. 调用 get_embedding 函数生成查询字符串的嵌入向量 query_embeddings
    query_embeddings = get_embeddings(query)
    # 2.文档列表documents生成向量，并存储在 documents_embeddings中
    documents_embeddings = embeddings.embed_documents(documents)
    # 3.计算查询嵌入向量与每个文档嵌入之间的余弦相似度，存储在similarities列表中
    similarities = [cosine_similarity(query_embeddings, doc_embedding) for doc_embedding in documents_embeddings]
    # 4 找到相似度最高的文档的索引 most_similar_index
    most_similar_index = similarities.index(max(similarities))
    return documents[most_similar_index], max(similarities)


if __name__ == '__main__':
    documents = [
        "OpenAI的ChatGPT是一个强大的语音模型。",
        "天空是蓝色的，阳光灿烂",
        "人工智能正在改变世界",
        "python是一种流行的编程语言"
    ]
    # query = "天空是什么颜色"
    query = "人工智能"
    most_similar_documents, similarity_score = search_documents(query, documents)
    print(f"最相似的文档：{most_similar_documents}")
    print(f"最相性得分：{similarity_score}")
